package txmusic;

import common.utils.TreeNode;

public class PruneLeaves {

    public TreeNode pruneLeaves (TreeNode root) {
        // write code here
        boolean flag = dfs(root);
        if (flag){
            return null;
        }
        return root;
    }

    public boolean dfs(TreeNode root){
        //做孩子是叶节点
        if (root.left != null && root.left.left == null && root.left.right == null){
            return true;
        }
        //右孩子是叶节点
        if (root.right != null && root.right.left == null && root.right.right == null){
            return true;
        }
        boolean left = false;
        boolean right = false;

        if (root.left != null){
            left = dfs(root.left);
        }
        if (root.right != null){
            right = dfs(root.right);
        }
        if (left){
            root.left = null;
        }
        if (right){
            root.right = null;
        }
        return false;
    }
}
